import argparse
import backtrader as bt
from datetime import datetime as dt
from time import time
from multiprocessing import cpu_count, Process
import os

from os.path import join, exists, dirname
if __name__ == '__main__':
    import sys
    sys.path.append(dirname(dirname(__file__)))
from strategies.netEasyTailwhip import StrategyClass
from utils.localdata import onebyone, download

params = dict(
        ma = 'S',         ## 'S': SMA、'E':EMA、'W':WMA
        period_b1 = 22,    # 前箱体长度
        period_b2 = 5,     # 后箱体长度
        boxR = 7,       # 前箱体检测，60、22、10 三均线最大差值放大boxR倍仍不高于股价
        breakdown = 2,  # 后箱体接受几次下突破
        needvcd = False,    # 需要量均死叉
        maxuplimit = -1,    # 后箱体等待期涨停上限，负数表示不限制
        postboxV = 1,      # 后箱体规则，1：公开课版本; 2: 正课版本，3：自动切换
        closeup1 = 1.05,   # 目标止盈
        closeup2 = 1.1,    # 顶部回落止盈
        closedown = 0.76,  # 止损
        code = None,
        name = None,
        log = True,
)

ps = dict(
        boxR = [7,7.4,7.8,8,8.2,8.6,9],       # 前箱体检测，60、22、10 三均线最大差值放大boxR倍仍不高于股价
        breakdown = [0,1,2],  # 后箱体接受几次下突破
        needvcd = [True,False],    # 需要量均死叉
        maxuplimit = [-1,0,1],    # 后箱体等待期涨停上限，负数表示不限制
        postboxV = [1,2],      # 后箱体规则，1：公开课版本; 2: 正课版本，3：自动切换
)

ps = dict(
        closeup1 = [x/100 for x in range(105, 118)],    # 目标止盈
        closeup2 = [x/100 for x in range(103, 110)],    # 顶部回落止盈
        closedown = [x/100 for x in range(72,90,2)],    # 止损
)
ps2 = None

def pps(p={}, ks=list(ps.keys()), ps=ps, ps2=ps2):
    if len(ks) > 0:
        for x in ps[ks[0]]:
            yield from pps(dict({f'{ks[0]}': x}, **p), ks[1:], ps, ps2)
    elif 'afterdown' in p and p['afterdown'] and ps2:
        yield from pps(p, list(ps2.keys()), ps2, None)
    else:
        yield p

print(len([x for x in pps()]))

def main(num, pss, StrategyClass=StrategyClass):
    logfile = join('logs', f's{num:03d}_{dt.now().isoformat()}_{os.getpid()}.log')
    def log(text, write=False, file=logfile, cache=dict(a=[])):
        lines = cache['a']
        lines.append(f'{dt.now().isoformat()} {text}')
        if write:
            with open(file, 'a', encoding='utf8') as f:
                f.write('\n'.join(lines))
            cache['a'] = ['\n']

    commparams = dict(log=False, closeup1=1.11, closeup2=1.06, closedown=0.8)
    commparams = {'log':False,'postboxV':1,'maxuplimit':-1,'needvcd':False,'breakdown':2,'boxR':7}
    import re
    replace = re.compile(r'0{4}0+1$')
    for p in pss:
        params = dict(commparams, **p)
        total = [0, 0, 0, 0, 0, 0] # 有交易的证券数，交易总笔数，盈利总笔数，损失总笔数，盈利总和, 证券总数
        # for daily_data, code, name in [d for d in onebyone(klen=300) if d[1] in ['sh.600378', 'sh.603333']]:
        for daily_data, code, name in onebyone(klen=300):
            params['code'] = code
            params['name'] = name
            cerebro = bt.Cerebro()
            # 设置启动资金
            cerebro.broker.setcash(10000000.0)
            # 设置佣金
            cerebro.broker.setcommission(commission=0.001)
            cerebro.addstrategy(StrategyClass, **params)
            data = bt.feeds.PandasData(dataname=daily_data)
            cerebro.adddata(data)
            cerebro.addsizer(bt.sizers.FixedSize, stake=100)  # 每笔交易使用固定交易量，stake代表交易量
            cerebro.addanalyzer(bt.analyzers.TradeAnalyzer, _name='ta')
            result = cerebro.run(plot=False)
            r = result[0].analyzers.ta.get_analysis()
            if r.total and r.total.total > 0:
                s = r.total.total
                l, w = 0, 0
                if r.total.total - r.total.open > 0:
                    l, w, x = r.lost.total, r.won.total, r.pnl.gross.total
                else:
                    x = cerebro.broker.getvalue() - 10000000
                log(f'{code} {name} {s} {w} {l} {x:.2f} ')
                total = [a+b for a,b in zip(total, [1,s,w,l,1,x])]
            else:
                total[-2] += 1
        total[-1] = round(total[-1], 2)
        log(replace.sub('', f'={total}={p}={dt.now()}'), True)

def parse_args():
    parser = argparse.ArgumentParser(
        description='神龙摆尾策略参数搜索-多进程')

    parser.add_argument('-n', default=cpu_count()-1, required=False, type=int, help='进程数')

    return parser.parse_args()

if __name__ == '__main__':
    args = parse_args()

    download()

    params = [x for x in pps()]
    size = (len(params) // (args.n - 1)) if args.n > 1 else len(params)

    for i in range(args.n):
        Process(target=main,kwargs={"num": i, "pss": params[size * i:size * (i + 1)]}).start()

    print(dt.now().isoformat(), f"主进程ID: {os.getpid()} 进程数量：{args.n} 参数组数量: {len(params)} 每个进程搜索参数组数量：{size}")
